use std::sync::Arc;

use axum::{
    async_trait,
    extract::{FromRequest, RawForm, Request},
    http::{
        header::{self, HeaderMap},
        HeaderValue, StatusCode,
    },
    response::{IntoResponse, Response},
    RequestExt,
};
use bytes::{BufMut, Bytes, BytesMut};
use serde::{de::DeserializeOwned, Serialize};
use tokio::sync::Mutex;
use validator::Validate;

use crate::error::AppError;

#[derive(Debug, Clone, Default)]
pub struct CollectBody(Arc<Mutex<Option<String>>>);

impl CollectBody {
    pub async fn set_body(self, body: String) {
        let mut s = self.0.lock().await;
        *s = Some(body);
    }

    pub async fn into_body(self) -> Option<String> {
        self.0.lock().await.take()
    }
}

pub struct Json<T>(pub T);

#[async_trait]
impl<T, S> FromRequest<S> for Json<T>
where
    T: DeserializeOwned + Send,

    S: Send + Sync,
{
    type Rejection = AppError;

    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let body = req.extensions().get::<CollectBody>().map(|f| (*f).clone());
        let (inner, bytes) = if json_content_type(req.headers()) {
            let bytes = Bytes::from_request(req, state)
                .await
                .map_err(|e| e.into_response())?;

            let inner = axum::Json::<T>::from_bytes(&bytes).map_err(|e| e.into_response())?;
            (inner.0, bytes)
        } else {
            match req.extract().await {
                Ok(RawForm(bytes)) => {
                    let value = serde_urlencoded::from_bytes::<T>(&bytes).map_err(|_| {
                        AppError::BadParam("Failed to deserialize form".to_string())
                    })?;
                    (value, bytes)
                }
                Err(err) => Err(err.into_response())?,
            }
        };

        if let Some(collect) = body {
            let ss = String::from_utf8_lossy(&bytes).to_string();
            collect.set_body(ss).await;
        }
        Ok(Self(inner))
    }
}

fn json_content_type(headers: &HeaderMap) -> bool {
    let content_type = if let Some(content_type) = headers.get(header::CONTENT_TYPE) {
        content_type
    } else {
        return false;
    };

    let content_type = if let Ok(content_type) = content_type.to_str() {
        content_type
    } else {
        return false;
    };

    let mime = if let Ok(mime) = content_type.parse::<mime::Mime>() {
        mime
    } else {
        return false;
    };

    let is_json_content_type = mime.type_() == "application"
        && (mime.subtype() == "json" || mime.suffix().map_or(false, |name| name == "json"));

    is_json_content_type
}

impl<T> IntoResponse for Json<T>
where
    T: Serialize,
{
    fn into_response(self) -> Response {
        // Use a small initial capacity of 128 bytes like serde_json::to_vec
        // https://docs.rs/serde_json/1.0.82/src/serde_json/ser.rs.html#2189
        let mut buf = BytesMut::with_capacity(128).writer();
        match serde_json::to_writer(&mut buf, &self.0) {
            Ok(()) => (
                [(
                    header::CONTENT_TYPE,
                    HeaderValue::from_static(mime::APPLICATION_JSON.as_ref()),
                )],
                buf.into_inner().freeze(),
            )
                .into_response(),
            Err(err) => (
                StatusCode::INTERNAL_SERVER_ERROR,
                [(
                    header::CONTENT_TYPE,
                    HeaderValue::from_static(mime::TEXT_PLAIN_UTF_8.as_ref()),
                )],
                err.to_string(),
            )
                .into_response(),
        }
    }
}

/// 带validate的json数据检测
pub struct ValJson<T>(pub T);

#[async_trait]
impl<T, S> FromRequest<S> for ValJson<T>
where
    T: DeserializeOwned + Validate + Send,

    S: Send + Sync,
{
    type Rejection = AppError;

    async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
        let inner: T = Json::from_request(req, state).await?.0;
        inner.validate()?;
        Ok(ValJson(inner))
    }
}
